#!/usr/bin/env python3
"""
代码沙箱测试基础类
提供通用的测试方法和工具函数
"""

import asyncio
import aiohttp
import json
import time
from typing import Dict, Any, List, Optional
from datetime import datetime

class BaseCodeSandboxTest:
    """代码沙箱测试基础类"""
    
    def __init__(self, base_url: str = "http://localhost:8000"):
        self.base_url = base_url
        self.session = None
        self.test_results = []
        self.code_type = "python"  # 默认代码类型，子类可覆盖
    
    async def __aenter__(self):
        """异步上下文管理器入口"""
        self.session = aiohttp.ClientSession()
        return self
    
    async def __aexit__(self, exc_type, exc_val, exc_tb):
        """异步上下文管理器出口"""
        if self.session:
            await self.session.close()
    
    async def check_service_health(self) -> bool:
        """检查服务健康状态"""
        try:
            if not self.session:
                self.session = aiohttp.ClientSession()
            
            async with self.session.get(f"{self.base_url}/health") as response:
                if response.status == 200:
                    return True
                else:
                    print(f"❌ 服务健康检查失败: HTTP {response.status}")
                    return False
        except Exception as e:
            print(f"❌ 无法连接到服务: {e}")
            print("请确保服务已启动 (python3 start_service.py)")
            return False
    
    async def submit_task(self, code: str, security_level: str = "low", 
                         timeout_seconds: int = 30, memory_limit_mb: int = 128,
                         cpu_limit: float = 0.5, environment_vars: Optional[Dict[str, str]] = None) -> str:
        """提交代码执行任务"""
        if not self.session:
            self.session = aiohttp.ClientSession()
            
        request_data = {
            "code": code,
            "code_type": self.code_type,
            "user_info": {
                "user_id": f"test_user_{int(time.time())}",
                "username": "测试用户",
                "permissions": ["execute"]
            },
            "security_level": security_level,
            "timeout_seconds": timeout_seconds,
            "memory_limit_mb": memory_limit_mb,
            "cpu_limit": cpu_limit,
            "environment_vars": environment_vars or {}
        }
        
        async with self.session.post(f"{self.base_url}/execute", json=request_data) as response:
            if response.status != 200:
                raise Exception(f"任务提交失败: HTTP {response.status}")
            
            result = await response.json()
            if not result.get("success"):
                raise Exception(f"任务提交失败: {result.get('message')}")
            
            return result["data"]["task_id"]
    
    async def wait_for_completion(self, task_id: str, max_wait: int = 60) -> Dict[str, Any]:
        """等待任务完成并返回结果"""
        if not self.session:
            self.session = aiohttp.ClientSession()
            
        start_time = time.time()
        
        while time.time() - start_time < max_wait:
            # 检查任务状态
            async with self.session.get(f"{self.base_url}/task/{task_id}/status") as response:
                if response.status != 200:
                    raise Exception(f"获取任务状态失败: HTTP {response.status}")
                
                status_data = await response.json()
                status = status_data["data"]["status"]
                
                if status in ["completed", "failed", "timeout"]:
                    # 获取完整结果
                    async with self.session.get(f"{self.base_url}/task/{task_id}/result") as result_response:
                        if result_response.status != 200:
                            raise Exception(f"获取任务结果失败: HTTP {result_response.status}")
                        
                        result_data = await result_response.json()
                        return result_data["data"]
                
                # 继续等待
                await asyncio.sleep(0.5)
        
        raise Exception(f"任务等待超时: {task_id}")
    
    async def run_test_case(self, test_case: Dict[str, Any]) -> Dict[str, Any]:
        """运行单个测试用例"""
        case_name = test_case.get("name", "未命名测试")
        print(f"  📋 测试用例: {case_name}")
        
        start_time = time.time()
        result = {
            "name": case_name,
            "start_time": start_time,
            "success": False,
            "error": None,
            "execution_result": None
        }
        
        try:
            # 提取测试参数
            code = test_case["code"]
            security_level = test_case.get("security_level", "low")
            timeout_seconds = test_case.get("timeout_seconds", 30)
            memory_limit_mb = test_case.get("memory_limit_mb", 128)
            cpu_limit = test_case.get("cpu_limit", 0.5)
            environment_vars = test_case.get("environment_vars", {})
            
            # 提交任务
            task_id = await self.submit_task(
                code=code,
                security_level=security_level,
                timeout_seconds=timeout_seconds,
                memory_limit_mb=memory_limit_mb,
                cpu_limit=cpu_limit,
                environment_vars=environment_vars
            )
            
            # 等待完成
            execution_result = await self.wait_for_completion(task_id, timeout_seconds + 10)
            result["execution_result"] = execution_result
            
            # 验证结果
            await self.validate_test_result(test_case, execution_result, result)
            
        except Exception as e:
            result["error"] = str(e)
            print(f"    ❌ 测试失败: {e}")
        
        result["end_time"] = time.time()
        result["duration"] = result["end_time"] - result["start_time"]
        self.test_results.append(result)
        
        if result["success"]:
            print(f"    ✅ 测试通过 ({result['duration']:.2f}s)")
        
        return result
    
    async def validate_test_result(self, test_case: Dict[str, Any], 
                                  execution_result: Dict[str, Any], 
                                  result: Dict[str, Any]):
        """验证测试结果"""
        expect_error = test_case.get("expect_error", False)
        expected_output = test_case.get("expected_output", [])
        
        # 检查执行状态
        status = execution_result.get("status")
        output = execution_result.get("output", "")
        error = execution_result.get("error", "")
        
        if expect_error:
            if status == "failed" or error:
                result["success"] = True
                print(f"    ✅ 预期错误已捕获: {status}")
            else:
                result["success"] = False
                print(f"    ❌ 期望出现错误但任务成功执行")
                return
        else:
            if status == "completed":
                result["success"] = True
                print(f"    ✅ 任务执行成功")
            elif status == "timeout":
                result["success"] = False
                print(f"    ❌ 任务执行超时")
                return
            else:
                result["success"] = False
                print(f"    ❌ 任务执行失败: {error}")
                return
        
        # 验证输出内容
        if expected_output and output:
            output_lines = [line.strip() for line in output.strip().split('\n') if line.strip()]
            
            for expected_line in expected_output:
                if not any(expected_line in line for line in output_lines):
                    result["success"] = False
                    print(f"    ❌ 期望输出未找到: '{expected_line}'")
                    print(f"    实际输出: {output_lines}")
                    return
            
            print(f"    ✅ 输出验证通过")
        
        # 验证执行时间
        execution_time = execution_result.get("execution_time")
        if execution_time:
            print(f"    ⏱️  执行时间: {execution_time:.2f}s")
        
        # 验证资源使用
        memory_usage = execution_result.get("memory_usage")
        if memory_usage:
            print(f"    💾 内存使用: {memory_usage}MB")
        
        cpu_usage = execution_result.get("cpu_usage")
        if cpu_usage:
            print(f"    🖥️  CPU使用: {cpu_usage:.1f}%")
    
    def print_test_summary(self):
        """打印测试总结"""
        total_tests = len(self.test_results)
        successful_tests = sum(1 for r in self.test_results if r["success"])
        failed_tests = total_tests - successful_tests
        
        print(f"\n📊 测试总结:")
        print(f"  总测试数: {total_tests}")
        print(f"  成功: {successful_tests} ✅")
        print(f"  失败: {failed_tests} ❌")
        print(f"  成功率: {(successful_tests/total_tests*100):.1f}%" if total_tests > 0 else "  成功率: 0%")
        
        if failed_tests > 0:
            print(f"\n❌ 失败的测试用例:")
            for result in self.test_results:
                if not result["success"]:
                    print(f"  - {result['name']}: {result.get('error', '验证失败')}")
        
        # 计算平均执行时间
        if self.test_results:
            avg_duration = sum(r.get("duration", 0) for r in self.test_results) / len(self.test_results)
            print(f"\n⏱️  平均测试时间: {avg_duration:.2f}s")
    
    async def test_service_status(self):
        """测试服务状态接口"""
        print("🔍 测试服务状态...")
        
        if not self.session:
            self.session = aiohttp.ClientSession()
        
        try:
            # 测试系统状态
            async with self.session.get(f"{self.base_url}/system/status") as response:
                data = await response.json()
                print(f"  系统状态: {data['success']}")
            
            # 测试队列信息
            async with self.session.get(f"{self.base_url}/queues/info") as response:
                data = await response.json()
                print(f"  队列信息: {data['success']}")
            
            # 测试执行器状态
            async with self.session.get(f"{self.base_url}/executors/status") as response:
                data = await response.json()
                print(f"  执行器状态: {data['success']}")
            
            print("  ✅ 服务状态检查完成")
        
        except Exception as e:
            print(f"  ❌ 服务状态检查失败: {e}") 